#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
基于 LangChain 的千问 Agent 最简实现
支持数学计算（Python代码执行）和天气查询（API调用）

使用方法:
1. 安装依赖: pip install -r requirements.txt
2. 设置环境变量: export DASHSCOPE_API_KEY='你的千问API_key'
3. 可选：设置天气API: export WEATHER_API_KEY='你的天气API_key'
4. 运行: python script_agent_02_agent_demo_cursor.py
"""

import os
import sys
import json
import requests
import traceback
from typing import Optional, Dict, Any, List
from io import StringIO
import contextlib

# LangChain imports
try:
    from langchain.agents import initialize_agent, AgentType, Tool
    from langchain.llms.base import LLM
    from langchain.callbacks.manager import CallbackManagerForLLMRun
    from langchain.schema import BaseMessage, HumanMessage, AIMessage
    from langchain.memory import ConversationBufferMemory
except ImportError:
    print("❌ 错误: 请先安装 LangChain")
    print("   运行: pip install langchain")
    sys.exit(1)


class QwenLLM(LLM):
    """千问 LLM 封装"""
    
    def __init__(self, api_key: str, model: str = "qwen-turbo"):
        self.api_key = api_key
        self.model = model
        self.url = "https://dashscope.aliyuncs.com/api/v1/services/aigc/text-generation/generation"
    
    @property
    def _llm_type(self) -> str:
        return "qwen"
    
    def _call(
        self,
        prompt: str,
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> str:
        """调用千问 API"""
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }
        
        data = {
            "model": self.model,
            "input": {
                "messages": [
                    {
                        "role": "user",
                        "content": prompt
                    }
                ]
            },
            "parameters": {
                "max_tokens": 2000,
                "temperature": 0.1
            }
        }
        
        try:
            response = requests.post(self.url, headers=headers, json=data, timeout=30)
            
            if response.status_code == 200:
                result = response.json()
                if 'output' in result and 'text' in result['output']:
                    return result['output']['text']
                else:
                    return f"API 响应格式异常: {result}"
            else:
                return f"API 请求失败 (状态码: {response.status_code}): {response.text}"
                
        except Exception as e:
            return f"调用千问 API 失败: {e}"


class MathCalculator:
    """数学计算工具 - 通过 Python 代码执行"""
    
    def __init__(self):
        # 安全的数学函数白名单
        self.safe_functions = {
            'abs', 'round', 'min', 'max', 'sum', 'len',
            'int', 'float', 'str', 'bool', 'list', 'tuple', 'dict', 'set',
            'range', 'enumerate', 'zip', 'sorted', 'reversed',
        }
        
        # 数学模块
        try:
            import math
            self.math_module = math
        except ImportError:
            self.math_module = None
    
    def calculate(self, expression: str) -> str:
        """安全执行数学计算"""
        try:
            # 基本安全检查
            dangerous_keywords = ['import', 'exec', 'eval', 'open', 'file', '__']
            for keyword in dangerous_keywords:
                if keyword in expression.lower():
                    return f"❌ 安全错误: 不允许使用 {keyword}"
            
            # 准备安全的执行环境
            safe_globals = {
                '__builtins__': {name: getattr(__builtins__, name) 
                               for name in self.safe_functions 
                               if hasattr(__builtins__, name)}
            }
            
            # 添加数学函数
            if self.math_module:
                for name in dir(self.math_module):
                    if not name.startswith('_'):
                        safe_globals[name] = getattr(self.math_module, name)
            
            # 捕获输出
            output = StringIO()
            with contextlib.redirect_stdout(output):
                # 执行代码
                result = eval(expression, safe_globals, {})
                
                # 如果有打印输出，包含在结果中
                printed_output = output.getvalue()
                if printed_output:
                    return f"计算结果: {result}\n输出: {printed_output.strip()}"
                else:
                    return f"计算结果: {result}"
                    
        except SyntaxError as e:
            return f"❌ 语法错误: {e}"
        except ZeroDivisionError:
            return "❌ 数学错误: 除零错误"
        except Exception as e:
            return f"❌ 计算错误: {e}"


class WeatherAPI:
    """天气查询工具"""
    
    def __init__(self, api_key: Optional[str] = None):
        self.api_key = api_key or os.getenv('WEATHER_API_KEY')
        # 使用免费的 OpenWeatherMap API
        self.base_url = "http://api.openweathermap.org/data/2.5/weather"
    
    def get_weather(self, city: str) -> str:
        """查询指定城市的天气"""
        if not self.api_key:
            # 如果没有 API key，返回模拟数据
            return self._mock_weather(city)
        
        try:
            params = {
                'q': city,
                'appid': self.api_key,
                'units': 'metric',  # 摄氏度
                'lang': 'zh_cn'     # 中文描述
            }
            
            response = requests.get(self.base_url, params=params, timeout=10)
            
            if response.status_code == 200:
                data = response.json()
                
                weather_info = {
                    '城市': data['name'],
                    '温度': f"{data['main']['temp']}°C",
                    '体感温度': f"{data['main']['feels_like']}°C",
                    '天气': data['weather'][0]['description'],
                    '湿度': f"{data['main']['humidity']}%",
                    '风速': f"{data['wind']['speed']} m/s"
                }
                
                result = f"🌤️ {city} 当前天气:\n"
                for key, value in weather_info.items():
                    result += f"   {key}: {value}\n"
                
                return result.strip()
                
            else:
                return f"❌ 天气查询失败: {response.text}"
                
        except Exception as e:
            return f"❌ 天气查询异常: {e}"
    
    def _mock_weather(self, city: str) -> str:
        """模拟天气数据（当没有 API key 时使用）"""
        return f"""🌤️ {city} 当前天气 (模拟数据):
   城市: {city}
   温度: 22°C
   体感温度: 24°C
   天气: 多云
   湿度: 65%
   风速: 3.2 m/s
   
💡 提示: 设置 WEATHER_API_KEY 环境变量获取真实天气数据"""


class AgentDemo:
    """Agent 主类"""
    
    def __init__(self):
        self.api_key = os.getenv('DASHSCOPE_API_KEY')
        if not self.api_key:
            raise ValueError("请设置环境变量 DASHSCOPE_API_KEY")
        
        # 初始化 LLM
        self.llm = QwenLLM(self.api_key)
        
        # 初始化工具
        self.math_calculator = MathCalculator()
        self.weather_api = WeatherAPI()
        
        # 定义工具
        self.tools = [
            Tool(
                name="数学计算",
                description="执行数学计算和Python表达式。输入应该是一个数学表达式，如 '2+3*4' 或 'math.sqrt(16)'",
                func=self.math_calculator.calculate
            ),
            Tool(
                name="天气查询",
                description="查询指定城市的当前天气信息。输入应该是城市名称，如 '北京' 或 '上海'",
                func=self.weather_api.get_weather
            )
        ]
        
        # 初始化记忆
        self.memory = ConversationBufferMemory(
            memory_key="chat_history",
            return_messages=True
        )
        
        # 初始化 agent
        self.agent = initialize_agent(
            tools=self.tools,
            llm=self.llm,
            agent=AgentType.CONVERSATIONAL_REACT_DESCRIPTION,
            memory=self.memory,
            verbose=True,
            max_iterations=3,
            early_stopping_method="generate"
        )
    
    def chat(self, message: str) -> str:
        """与 agent 对话"""
        try:
            response = self.agent.run(input=message)
            return response
        except Exception as e:
            error_msg = f"Agent 执行错误: {e}"
            print(f"❌ {error_msg}")
            traceback.print_exc()
            return error_msg
    
    def add_tool(self, tool: Tool):
        """动态添加工具（扩展性支持）"""
        self.tools.append(tool)
        # 重新初始化 agent
        self.agent = initialize_agent(
            tools=self.tools,
            llm=self.llm,
            agent=AgentType.CONVERSATIONAL_REACT_DESCRIPTION,
            memory=self.memory,
            verbose=True,
            max_iterations=3,
            early_stopping_method="generate"
        )
        print(f"✅ 已添加工具: {tool.name}")


def main():
    """主函数"""
    print("=" * 60)
    print("🤖 千问 Agent Demo - 支持数学计算和天气查询")
    print("=" * 60)
    
    try:
        # 初始化 agent
        agent = AgentDemo()
        print("✅ Agent 初始化成功!")
        print("\n💡 支持的功能:")
        print("   📊 数学计算: 如 '计算 2+3*4' 或 '求平方根 16'")
        print("   🌤️  天气查询: 如 '北京天气' 或 '上海的天气怎么样'")
        print("   💬 普通对话: 任何其他问题")
        print("\n输入 'quit' 或 'exit' 退出")
        print("-" * 60)
        
        while True:
            try:
                user_input = input("\n👤 你: ").strip()
                
                if user_input.lower() in ['quit', 'exit', '退出']:
                    print("👋 再见!")
                    break
                
                if not user_input:
                    continue
                
                print("\n🤖 Agent:", end=" ")
                response = agent.chat(user_input)
                print(response)
                print("-" * 60)
                
            except KeyboardInterrupt:
                print("\n👋 再见!")
                break
            except Exception as e:
                print(f"\n❌ 错误: {e}")
                continue
                
    except Exception as e:
        print(f"❌ 初始化失败: {e}")
        print("\n📝 检查事项:")
        print("   1. 确认已安装所有依赖: pip install -r requirements.txt")
        print("   2. 确认已设置 DASHSCOPE_API_KEY 环境变量")
        print("   3. 确认网络连接正常")


if __name__ == "__main__":
    main()
